import os
import pandas as pd
from src import log
import chardet


class Files:

    @staticmethod
    def get_content(file_name, index):
        _content = pd.DataFrame()
        # 判断文件格式并读取
        file_ext = os.path.splitext(file_name)
        if file_ext[1] in [".xls", ".xlsx"]:
            _content = pd.read_excel(file_name)
        elif file_ext[1] == ".csv":
            f = open(file_name, "rb")
            encode = chardet.detect(f.read())["encoding"]
            if encode == 'GB2312':
                try:
                    _content = pd.read_csv(file_name, encoding="GBK")
                except UnicodeDecodeError:
                    _content = pd.read_csv(file_name, encoding="GB18030")
            elif encode == 'UTF-8-SIG':
                _content = pd.read_csv(file_name, encoding="UTF-8-SIG'")
            elif encode == 'iso-8859-1':
                _content = pd.read_csv(file_name, encoding="gbk")

        # python选取特定列——pandas的loc使用（列切片及行切片）
        content = _content.loc[:, :] if index == [''] else _content.loc[:,index]

        return content

    @staticmethod
    def write_file_sheet(file_name, result, sheet):
        if os.path.exists(file_name):
            writer = pd.ExcelWriter(file_name, mode='a', engine='openpyxl')
            result.to_excel(writer, sheet_name=sheet)
            writer.close()
        else:
            result.to_excel(file_name, sheet_name=sheet)
            writer = pd.ExcelWriter(file_name, mode='a',
                                    engine='openpyxl')  # 追加sheet页
            writer.close()

        log.info("Write Succeed")

    @staticmethod
    def read_excel(path, sheet_no, page_num, page_size):
        # 列名及总条数
        df = pd.read_excel(path, sheet_name=sheet_no)
        df_columns = df.columns.values.tolist()
        df_total = df.shape[0]
        # 分页数据
        df = pd.read_excel(path, sheet_name=sheet_no, nrows=page_size, skiprows=(page_num - 1) * page_size,
                           keep_default_na=False)
        df.columns = df_columns
        df_list = df.to_dict(orient='records')
        return df_columns, df_total, df_list

